
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import os

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.mean = torch.as_tensor((0.485, 0.456, 0.406), dtype=torch.float32)
        self.std  = torch.as_tensor((0.229, 0.224, 0.225), dtype=torch.float32)
        self.max  = torch.as_tensor((255.0), dtype=torch.float32)

    def forward(self, x):
        x = x.permute(0, 2, 3, 1)
        x = (x / self.max - self.mean) / self.std
        x = x.permute(0, 3, 1, 2)
        return x
def genProModel():

    model = Model()
    input = torch.ones(1, 3, 960, 960)  # chw  rgb
    with torch.no_grad():
        y = model(input)
    print(y.shape)

    torch.onnx.export(
        model, (input,), 
        # 储存的文件路径

        "normal.onnx",  
        # 为输入和输出节点指定名称，方便后面查看或者操作
        input_names=["images"], 
        output_names=["output"], 
        opset_version=12, 
        # 表示他有batch、height、width3个维度是动态的，在onnx中给其赋值为-1
        # 通常，我们只设置batch为动态，其他的避免动态
        dynamic_axes={
            "images": {0: "batch"},
            "output": {0: "batch"},
        }
    )



if __name__ == "__main__":
    genProModel()
